#include <stdio.h>
#include <sys/socket.h>
#include <stdlib.h>
#include <sys/types.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <string.h>
#include <netinet/in.h>
#include <sys/stat.h>
#include <pthread.h>
#include <errno.h>

#define BUFFER_SIZE 1024    // 数据缓冲区大小
#define FILE_CHUNK_SIZE 512 // 文件分片大小
#define MAX_FILENAME 256    // 最大文件名长度

// 定义消息类型
typedef enum {
    MSG_TEXT = 1,        // 普通文本消息
    MSG_FILE_SEND,       // 发送文件请求
    MSG_FILE_RECV,       // 接收文件请求
    MSG_FILE_DATA,       // 文件数据包
    MSG_FILE_END,        // 文件传输结束
    MSG_ACK,             // 确认消息
    MSG_ERROR            // 错误消息
} msg_type_t;

// 消息头结构体
typedef struct {
    msg_type_t type;           // 消息类型
    int seq_num;               // 序列号（用于文件分片）
    int data_len;              // 数据长度
    char filename[MAX_FILENAME]; // 文件名
} msg_header_t;

// 完整消息结构体
typedef struct {
    msg_header_t header;
    char data[FILE_CHUNK_SIZE];
} message_t;

// 线程参数结构体
typedef struct {
    int sockfd;
    char filename[MAX_FILENAME];
    int is_send;  // 1表示发送文件，0表示接收文件
} thread_param_t;

// TCP发送函数 - 确保发送完整数据
ssize_t tcp_send_all(int sockfd, const void *buf, size_t len) {
    size_t total_sent = 0;
    const char *ptr = (const char *)buf;
    
    while (total_sent < len) {
        ssize_t sent = send(sockfd, ptr + total_sent, len - total_sent, 0);
        if (sent <= 0) {
            if (sent < 0) {
                perror("TCP发送失败");
            }
            return -1;
        }
        total_sent += sent;
    }
    return total_sent;
}

// TCP接收函数 - 确保接收完整数据
ssize_t tcp_recv_all(int sockfd, void *buf, size_t len) {
    size_t total_recv = 0;
    char *ptr = (char *)buf;
    
    while (total_recv < len) {
        ssize_t recvd = recv(sockfd, ptr + total_recv, len - total_recv, 0);
        if (recvd <= 0) {
            if (recvd < 0) {
                perror("TCP接收失败");
            }
            return -1;
        }
        total_recv += recvd;
    }
    return total_recv;
}

// 发送文件线程函数
void* send_file_thread(void *arg) {
    thread_param_t *param = (thread_param_t *)arg;
    int sockfd = param->sockfd;
    const char *filename = param->filename;
    
    FILE *fp = fopen(filename, "rb");
    if (fp == NULL) {
        printf("无法打开文件: %s\n", filename);
        free(param);
        return NULL;
    }

    // 获取文件大小
    fseek(fp, 0, SEEK_END);
    long file_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    printf("开始发送文件: %s (大小: %ld 字节)\n", filename, file_size);

    message_t msg;
    int seq_num = 0;
    int bytes_read;
    
    // 发送文件数据
    while ((bytes_read = fread(msg.data, 1, FILE_CHUNK_SIZE, fp)) > 0) {
        // 设置消息头
        msg.header.type = MSG_FILE_DATA;
        msg.header.seq_num = seq_num++;
        msg.header.data_len = bytes_read;
        strncpy(msg.header.filename, filename, MAX_FILENAME - 1);
        msg.header.filename[MAX_FILENAME - 1] = '\0';

        // 发送数据包
        if (tcp_send_all(sockfd, &msg, sizeof(msg_header_t) + bytes_read) == -1) {
            printf("发送文件数据失败\n");
            fclose(fp);
            free(param);
            return NULL;
        }

        // 等待ACK确认
        message_t ack_msg;
        if (tcp_recv_all(sockfd, &ack_msg, sizeof(msg_header_t)) == -1) {
            printf("接收ACK失败\n");
            fclose(fp);
            free(param);
            return NULL;
        }

        if (ack_msg.header.type != MSG_ACK || ack_msg.header.seq_num != seq_num - 1) {
            printf("收到错误的ACK\n");
            fclose(fp);
            free(param);
            return NULL;
        }

        printf("已发送第 %d 个数据包 (%d 字节)\n", seq_num - 1, bytes_read);
    }

    // 发送文件结束标志
    msg.header.type = MSG_FILE_END;
    msg.header.seq_num = seq_num;
    msg.header.data_len = 0;
    strncpy(msg.header.filename, filename, MAX_FILENAME - 1);
    
    if (tcp_send_all(sockfd, &msg, sizeof(msg_header_t)) == -1) {
        printf("发送文件结束标志失败\n");
        fclose(fp);
        free(param);
        return NULL;
    }

    // 等待最后的ACK
    message_t final_ack;
    if (tcp_recv_all(sockfd, &final_ack, sizeof(msg_header_t)) == -1) {
        printf("接收最终ACK失败\n");
    }

    fclose(fp);
    printf("文件发送完成: %s\n", filename);
    free(param);
    return NULL;
}

// 接收文件线程函数
void* receive_file_thread(void *arg) {
    thread_param_t *param = (thread_param_t *)arg;
    int sockfd = param->sockfd;
    const char *filename = param->filename;
    
    printf("[接收线程] 开始请求文件: %s\n", filename);
    
    message_t request_msg;
    
    // 发送接收文件请求
    request_msg.header.type = MSG_FILE_RECV;
    request_msg.header.seq_num = 0;
    request_msg.header.data_len = 0;
    strncpy(request_msg.header.filename, filename, MAX_FILENAME - 1);
    request_msg.header.filename[MAX_FILENAME - 1] = '\0';

    printf("[接收线程] 发送文件请求: %s\n", filename);
    if (tcp_send_all(sockfd, &request_msg, sizeof(msg_header_t)) == -1) {
        printf("[接收线程] 发送文件请求失败\n");
        free(param);
        return NULL;
    }

    // 创建本地文件（添加_received后缀避免覆盖）
    char local_filename[MAX_FILENAME + 20];
    snprintf(local_filename, sizeof(local_filename), "%s_received", filename);
    
    FILE *fp = fopen(local_filename, "wb");
    if (fp == NULL) {
        printf("[接收线程] 无法创建文件: %s (错误: %s)\n", local_filename, strerror(errno));
        free(param);
        return NULL;
    }

    printf("[接收线程] 等待服务器响应...\n");

    message_t msg, ack_msg;
    int expected_seq = 0;
    int total_bytes = 0;

    while (1) {
        // 先接收消息头
        printf("[接收线程] 等待接收消息头...\n");
        ssize_t header_recv = tcp_recv_all(sockfd, &msg.header, sizeof(msg_header_t));
        if (header_recv == -1) {
            printf("[接收线程] 接收文件消息头失败\n");
            fclose(fp);
            unlink(local_filename);
            free(param);
            return NULL;
        }

        printf("[接收线程] 收到消息类型: %d, 序列号: %d, 数据长度: %d\n", 
               msg.header.type, msg.header.seq_num, msg.header.data_len);

        // 处理不同类型的消息
        if (msg.header.type == MSG_FILE_DATA) {
            // 检查序列号
            if (msg.header.seq_num != expected_seq) {
                printf("[接收线程] 收到乱序的数据包，期望: %d，实际: %d\n", 
                       expected_seq, msg.header.seq_num);
                
                // 发送ACK（即使是乱序包也要确认）
                ack_msg.header.type = MSG_ACK;
                ack_msg.header.seq_num = msg.header.seq_num;
                ack_msg.header.data_len = 0;
                tcp_send_all(sockfd, &ack_msg, sizeof(msg_header_t));
                continue;
            }

            // 接收文件数据
            if (msg.header.data_len > 0) {
                printf("[接收线程] 接收文件数据 %d 字节...\n", msg.header.data_len);
                if (tcp_recv_all(sockfd, msg.data, msg.header.data_len) == -1) {
                    printf("[接收线程] 接收文件数据失败\n");
                    fclose(fp);
                    unlink(local_filename);
                    free(param);
                    return NULL;
                }

                // 写入文件数据
                if (fwrite(msg.data, 1, msg.header.data_len, fp) != msg.header.data_len) {
                    printf("[接收线程] 写入文件失败: %s\n", strerror(errno));
                    fclose(fp);
                    unlink(local_filename);
                    free(param);
                    return NULL;
                }

                total_bytes += msg.header.data_len;
                expected_seq++;
                printf("[接收线程] 接收到第 %d 个数据包 (%d 字节)，总计: %d 字节\n", 
                       msg.header.seq_num, msg.header.data_len, total_bytes);
            }

            // 发送ACK
            ack_msg.header.type = MSG_ACK;
            ack_msg.header.seq_num = msg.header.seq_num;
            ack_msg.header.data_len = 0;
            
            if (tcp_send_all(sockfd, &ack_msg, sizeof(msg_header_t)) == -1) {
                printf("[接收线程] 发送ACK失败\n");
            } else {
                printf("[接收线程] 已发送ACK for seq %d\n", msg.header.seq_num);
            }

        } else if (msg.header.type == MSG_FILE_END) {
            printf("[接收线程] 收到文件结束标志\n");
            
            // 发送最终ACK
            ack_msg.header.type = MSG_ACK;
            ack_msg.header.seq_num = msg.header.seq_num;
            ack_msg.header.data_len = 0;
            tcp_send_all(sockfd, &ack_msg, sizeof(msg_header_t));
            
            printf("[接收线程] 文件接收完成: %s (总大小: %d 字节)\n", local_filename, total_bytes);
            break;
            
        } else if (msg.header.type == MSG_ERROR) {
            printf("[接收线程] 服务器错误: 文件 %s 不存在或无法读取\n", filename);
            fclose(fp);
            unlink(local_filename);
            free(param);
            return NULL;
        } else {
            printf("[接收线程] 收到未知消息类型: %d\n", msg.header.type);
        }
    }

    fclose(fp);
    free(param);
    printf("[接收线程] 线程结束\n");
    return NULL;
}

// 显示帮助信息
void show_help() {
    printf("\n可用命令:\n");
    printf("  send <filename>  - 发送文件到服务器\n");
    printf("  get <filename>   - 从服务器获取文件\n");
    printf("  help            - 显示帮助信息\n");
    printf("  quit            - 退出程序\n");
    printf("  其他文本         - 发送普通消息\n\n");
}

int main(int argc, char const *argv[]) {
    if (argc < 3) {
        fprintf(stderr, "用法: %s <服务器IP> <端口>\n", argv[0]);
        exit(1);
    }

    // 1.创建套接字
    int sockfd;
    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
        perror("创建套接字失败");
        exit(1);
    }

    printf("套接字创建成功，sockfd = %d\n", sockfd);

    // 2.服务器地址信息结构体
    struct sockaddr_in serveraddr;
    serveraddr.sin_family = AF_INET;
    serveraddr.sin_addr.s_addr = inet_addr(argv[1]);
    serveraddr.sin_port = htons(atoi(argv[2]));

    // 3.连接到服务器
    if (connect(sockfd, (struct sockaddr *)&serveraddr, sizeof(serveraddr)) == -1) {
        perror("连接服务器失败");
        close(sockfd);
        exit(1);
    }

    printf("成功连接到服务器 %s:%s\n", argv[1], argv[2]);
    show_help();

    // 4.主循环 - 处理用户输入
    char input[BUFFER_SIZE];
    while (1) {
        printf("客户端> ");
        fflush(stdout);
        
        if (fgets(input, BUFFER_SIZE, stdin) == NULL) {
            break;
        }

        // 移除换行符
        input[strcspn(input, "\n")] = '\0';
        
        if (strlen(input) == 0) {
            continue;
        }

        // 解析命令
        char command[32], filename[MAX_FILENAME];
        if (sscanf(input, "%s %s", command, filename) == 2) {
            if (strcmp(command, "send") == 0) {
                // 创建发送文件线程
                thread_param_t *param = malloc(sizeof(thread_param_t));
                param->sockfd = sockfd;
                strcpy(param->filename, filename);
                param->is_send = 1;
                
                pthread_t thread_id;
                if (pthread_create(&thread_id, NULL, send_file_thread, param) != 0) {
                    perror("创建发送文件线程失败");
                    free(param);
                } else {
                    pthread_detach(thread_id);  // 分离线程，自动回收资源
                }
                continue;
            } else if (strcmp(command, "get") == 0) {
                // 创建接收文件线程
                thread_param_t *param = malloc(sizeof(thread_param_t));
                param->sockfd = sockfd;
                strcpy(param->filename, filename);
                param->is_send = 0;
                
                pthread_t thread_id;
                if (pthread_create(&thread_id, NULL, receive_file_thread, param) != 0) {
                    perror("创建接收文件线程失败");
                    free(param);
                } else {
                    pthread_detach(thread_id);  // 分离线程，自动回收资源
                }
                continue;
            }
        }
        
        // 处理单个命令
        if (strcmp(input, "quit") == 0) {
            break;
        } else if (strcmp(input, "help") == 0) {
            show_help();
            continue;
        }

        // 发送普通文本消息
        message_t msg;
        msg.header.type = MSG_TEXT;
        msg.header.seq_num = 0;
        msg.header.data_len = strlen(input);
        memset(msg.header.filename, 0, MAX_FILENAME);
        strcpy(msg.data, input);

        if (tcp_send_all(sockfd, &msg, sizeof(msg_header_t) + msg.header.data_len) == -1) {
            printf("发送消息失败\n");
            continue;
        }

        // 接收服务器回复
        message_t reply;
        if (tcp_recv_all(sockfd, &reply.header, sizeof(msg_header_t)) == -1) {
            printf("接收回复失败\n");
            continue;
        }

        if (reply.header.data_len > 0) {
            if (tcp_recv_all(sockfd, reply.data, reply.header.data_len) == -1) {
                printf("接收回复数据失败\n");
                continue;
            }
            reply.data[reply.header.data_len] = '\0';
            printf("服务器回复: %s\n", reply.data);
        }
    }

    // 5.关闭套接字
    close(sockfd);
    printf("客户端已退出\n");
    return 0;
}